#!/usr/bin/env python3
import subprocess
import psutil
import time
import logging
from datetime import datetime
import os
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.utils import formataddr

class LinuxDefender:
    def __init__(self):
        self.logger = self._setup_logging()
        self.connection_limit = 100  # 每个IP的最大连接数
        self.banned_ips = set()
        self.whitelist_ips = set()  # 新增白名单
        self.alert_settings = {
            'cpu_threshold': 80,
            'memory_threshold': 80,
            'connection_threshold': 1000,
            'email_notification': False,
            'email_address': '',
            'smtp_server': 'smtp.qq.com',  # 默认使用QQ邮箱
            'smtp_port': 587,
            'smtp_username': '',
            'smtp_password': ''  # 邮箱授权码
        }
        self.last_alert_time = {}  # 用于控制告警频率
        
        # 添加历史数据存储
        self.history = {
            'timestamps': [],
            'cpu': [],
            'memory': [],
            'network_in': [],
            'network_out': [],
            'connections': [],
            'disk_read': [],
            'disk_write': []
        }
        self.max_history_points = 60  # 保存最近60个数据点
        self.last_network_bytes = None  # 用于计算网络速率
        
        # 修改缓存设置
        self.stats_cache = None
        self.last_stats_update = 0
        self.stats_cache_ttl = 5  # 缓存有效期改为5秒
        self.stats_updating = False  # 添加更新锁
        
        # 安全防御设置
        self.security_settings = {
            'ip_defense': {
                'max_failed_attempts': 5,  # 最大失败尝试次数
                'block_time': 3600,        # 封禁时间（秒）
                'suspicious_ports': [22, 3389, 8080],  # 敏感端口
                'whitelist_countries': ['CN'],  # 允许访问的国家/地区
            },
            'process_defense': {
                'max_cpu_percent': 90,     # CPU使用率阈值
                'max_memory_percent': 90,   # 内存使用率阈值
                'suspicious_names': ['cryptominer', 'malware'],  # 可疑进程名
                'protected_processes': ['sshd', 'nginx'],  # 受保护进程
            },
            'file_defense': {
                'protected_paths': ['/etc/passwd', '/etc/shadow'],  # 受保护文件
                'monitored_dirs': ['/var/www', '/etc'],  # 监控目录
                'file_types': ['.php', '.sh', '.py'],    # 监控文件类型
            },
            'network_defense': {
                'max_connections': 1000,    # 最大连接数
                'rate_limit': 100,         # 每秒请求限制
                'ddos_threshold': 1000,    # DDoS检测阈值
            }
        }
        
        # 初始化防御组件
        self.failed_attempts = {}      # 记录登录失败
        self.file_checksums = {}       # 文件校验和
        self.network_stats = {}        # 网络统计
        self.process_history = {}      # 进程历史
        
    def _setup_logging(self):
        try:
            logging.basicConfig(
                filename='defender.log',
                level=logging.INFO,
                format='%(asctime)s - %(levelname)s - %(message)s'
            )
            logger = logging.getLogger('LinuxDefender')
            
            # 添加控制台输出
            console_handler = logging.StreamHandler()
            console_handler.setLevel(logging.INFO)
            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
            console_handler.setFormatter(formatter)
            logger.addHandler(console_handler)
            
            return logger
        except Exception as e:
            print(f"日志设置失败: {str(e)}")
            return logging.getLogger('LinuxDefender')

    def is_ip_whitelisted(self, ip):
        """检查IP是否在白名单中"""
        return ip in self.whitelist_ips
        
    def add_to_whitelist(self, ip):
        """添加IP到白名单"""
        self.whitelist_ips.add(ip)
        if ip in self.banned_ips:
            self.unban_ip(ip)
        self.logger.info(f"已添加IP到白名单: {ip}")
        
    def remove_from_whitelist(self, ip):
        """从白名单移除IP"""
        self.whitelist_ips.remove(ip)
        self.logger.info(f"已从白名单移除IP: {ip}")

    def check_connections(self):
        """检查并统计当前所有TCP连接"""
        connections = psutil.net_connections(kind='tcp')
        ip_count = {}
        
        for conn in connections:
            if conn.raddr and conn.raddr.ip:
                ip = conn.raddr.ip
                if self.is_ip_whitelisted(ip):
                    continue
                    
                ip_count[ip] = ip_count.get(ip, 0) + 1
                
                if ip_count[ip] > self.connection_limit and ip not in self.banned_ips:
                    self.ban_ip(ip)

    def ban_ip(self, ip):
        """使用iptables封禁IP"""
        try:
            subprocess.run([
                'iptables', 
                '-A', 
                'INPUT', 
                '-s', 
                ip, 
                '-j', 
                'DROP'
            ], check=True)
            self.banned_ips.add(ip)
            self.logger.info(f"已封禁IP: {ip}")
        except subprocess.CalledProcessError as e:
            self.logger.error(f"封禁IP失败: {ip}, 错误: {str(e)}")

    def unban_ip(self, ip):
        """解除IP封禁"""
        try:
            subprocess.run([
                'iptables',
                '-D',
                'INPUT',
                '-s',
                ip,
                '-j',
                'DROP'
            ], check=True)
            self.banned_ips.remove(ip)
            self.logger.info(f"已解除封禁IP: {ip}")
        except subprocess.CalledProcessError as e:
            self.logger.error(f"解除封禁IP失败: {ip}, 错误: {str(e)}")

    def get_recent_logs(self, lines=100):
        """获取最近的日志记录"""
        try:
            if not os.path.exists('defender.log'):
                return []
                
            with open('defender.log', 'r') as f:
                logs = f.readlines()[-lines:]
            return [{'timestamp': log.split(' - ')[0],
                    'level': log.split(' - ')[1],
                    'message': ' - '.join(log.split(' - ')[2:]).strip()} 
                   for log in logs]
        except Exception as e:
            self.logger.error(f"获取日志失败: {str(e)}")
            return []

    def update_alert_settings(self, settings):
        """更新告警设置"""
        self.alert_settings.update(settings)
        self.logger.info(f"已更新告警设置: {settings}")

    def send_alert_email(self, subject, message):
        """发送告警邮件"""
        try:
            if not self.alert_settings['email_notification'] or not self.alert_settings['email_address']:
                return False
                
            # 检查告警频率（同一类型的告警间隔至少5分钟）
            current_time = time.time()
            if subject in self.last_alert_time:
                if current_time - self.last_alert_time[subject] < 300:  # 5分钟
                    return False
            
            msg = MIMEMultipart()
            msg['From'] = formataddr(("Linux防护系统", self.alert_settings['smtp_username']))
            msg['To'] = self.alert_settings['email_address']
            msg['Subject'] = subject
            
            # 添加HTML内容
            html_content = f"""
            <html>
                <body>
                    <h2>{subject}</h2>
                    <p>{message}</p>
                    <hr>
                    <p>系统状态:</p>
                    <ul>
                        <li>CPU使用率: {psutil.cpu_percent()}%</li>
                        <li>内存使用率: {psutil.virtual_memory().percent}%</li>
                        <li>时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</li>
                    </ul>
                </body>
            </html>
            """
            msg.attach(MIMEText(html_content, 'html'))
            
            # 连接SMTP服务器并发送
            with smtplib.SMTP(self.alert_settings['smtp_server'], self.alert_settings['smtp_port']) as server:
                server.starttls()
                server.login(self.alert_settings['smtp_username'], 
                           self.alert_settings['smtp_password'])
                server.send_message(msg)
            
            self.last_alert_time[subject] = current_time
            self.logger.info(f"告警邮件已发送: {subject}")
            return True
            
        except Exception as e:
            self.logger.error(f"发送告警邮件失败: {str(e)}")
            return False

    def check_and_alert(self):
        """检查系统状态并发送告警"""
        cpu_percent = psutil.cpu_percent()
        memory_percent = psutil.virtual_memory().percent
        connections = len([c for c in psutil.net_connections() if c.status == 'ESTABLISHED'])
        
        alerts = []
        
        if cpu_percent > self.alert_settings['cpu_threshold']:
            alerts.append(f"CPU使用率 ({cpu_percent}%) 超过阈值 ({self.alert_settings['cpu_threshold']}%)")
            
        if memory_percent > self.alert_settings['memory_threshold']:
            alerts.append(f"内存使用率 ({memory_percent}%) 超过阈值 ({self.alert_settings['memory_threshold']}%)")
            
        if connections > self.alert_settings['connection_threshold']:
            alerts.append(f"连接数 ({connections}) 超过阈值 ({self.alert_settings['connection_threshold']})")
            
        if alerts:
            subject = "Linux防护系统 - 系统资源告警"
            message = "\n".join(alerts)
            self.send_alert_email(subject, message)

    def get_system_stats(self):
        """获取系统状态（带缓存）"""
        current_time = time.time()
        
        # 如果缓存有效，直接返回缓存数据
        if (self.stats_cache is not None and 
            current_time - self.last_stats_update < self.stats_cache_ttl):
            return self.stats_cache
            
        # 如果正在更新，返回旧的缓存数据
        if self.stats_updating and self.stats_cache is not None:
            return self.stats_cache
            
        try:
            self.stats_updating = True
            
            # CPU信息
            cpu_percent = psutil.cpu_percent(interval=0.1)
            try:
                cpu_freq = psutil.cpu_freq()
                if cpu_freq:
                    cpu_frequency = {
                        'current': cpu_freq.current,
                        'min': cpu_freq.min,
                        'max': cpu_freq.max
                    }
                else:
                    cpu_frequency = {
                        'current': 0,
                        'min': 0,
                        'max': 0
                    }
            except (AttributeError, OSError):
                cpu_frequency = {
                    'current': 0,
                    'min': 0,
                    'max': 0
                }
            
            # 系统负载
            try:
                load1, load5, load15 = os.getloadavg()
                load_average = {
                    '1min': load1,
                    '5min': load5,
                    '15min': load15
                }
            except (AttributeError, OSError):
                load_average = {
                    '1min': 0,
                    '5min': 0,
                    '15min': 0
                }
            
            # 内存信息
            memory = psutil.virtual_memory()
            try:
                swap = psutil.swap_memory()
                swap_info = {
                    'total': swap.total,
                    'used': swap.used,
                    'free': swap.free,
                    'percent': swap.percent
                }
            except (AttributeError, OSError):
                swap_info = {
                    'total': 0,
                    'used': 0,
                    'free': 0,
                    'percent': 0
                }
            
            # 磁盘信息 - 只获取根分区
            disk = psutil.disk_usage('/')
            disk_io = psutil.disk_io_counters()
            
            # 网络信息 - 简化数据
            net_io = psutil.net_io_counters()
            
            # 获取活跃连接数而不是详细连接信息
            connections = psutil.net_connections()
            conn_stats = {
                'ESTABLISHED': len([c for c in connections if c.status == 'ESTABLISHED']),
                'TIME_WAIT': len([c for c in connections if c.status == 'TIME_WAIT']),
                'CLOSE_WAIT': len([c for c in connections if c.status == 'CLOSE_WAIT']),
                'LISTEN': len([c for c in connections if c.status == 'LISTEN'])
            }
            
            # 网络接口信息 - 只获取主要信息
            network_interfaces = {}
            for name, addrs in psutil.net_if_addrs().items():
                network_interfaces[name] = {
                    'addresses': [addr.address for addr in addrs if addr.family == 2][:1]  # 只取第一个IPv4地址
                }
            
            # 进程信息 - 只获取总数和前5个占用最多资源的进程
            processes = []
            for proc in sorted(psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']), 
                             key=lambda p: p.info['cpu_percent'] or 0, 
                             reverse=True)[:5]:
                try:
                    processes.append({
                        'pid': proc.info['pid'],
                        'name': proc.info['name'],
                        'cpu_percent': proc.info['cpu_percent'],
                        'memory_percent': proc.info['memory_percent']
                    })
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            
            # 构建状态数据
            stats = {
                'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'cpu': {
                    'percent': cpu_percent,
                    'cores': psutil.cpu_count(),
                    'frequency': cpu_frequency,
                    'load_average': load_average
                },
                'memory': {
                    'total': memory.total,
                    'used': memory.used,
                    'free': memory.free,
                    'percent': memory.percent,
                    'swap': swap_info
                },
                'disk': {
                    'total': disk.total,
                    'used': disk.used,
                    'free': disk.free,
                    'percent': disk.percent,
                    'io': {
                        'read_bytes': disk_io.read_bytes,
                        'write_bytes': disk_io.write_bytes
                    }
                },
                'network': {
                    'bytes_sent': net_io.bytes_sent,
                    'bytes_recv': net_io.bytes_recv,
                    'connections': conn_stats,
                    'interfaces': network_interfaces
                },
                'processes': {
                    'total': len(processes),
                    'top': processes
                }
            }
            
            # 更新缓存
            self.stats_cache = stats
            self.last_stats_update = current_time
            
            return stats
            
        except Exception as e:
            self.logger.error(f"获取系统状态失败: {str(e)}")
            return self.stats_cache or {
                'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'cpu': {
                    'percent': 0,
                    'cores': 1,
                    'frequency': {'current': 0, 'min': 0, 'max': 0},
                    'load_average': {'1min': 0, '5min': 0, '15min': 0}
                },
                'memory': {
                    'total': 0,
                    'used': 0,
                    'free': 0,
                    'percent': 0,
                    'swap': {
                        'total': 0,
                        'used': 0,
                        'free': 0,
                        'percent': 0
                    }
                },
                'disk': {
                    'total': 0,
                    'used': 0,
                    'free': 0,
                    'percent': 0,
                    'io': {
                        'read_bytes': 0,
                        'write_bytes': 0
                    }
                },
                'network': {
                    'bytes_sent': 0,
                    'bytes_recv': 0,
                    'connections': {
                        'ESTABLISHED': 0,
                        'TIME_WAIT': 0,
                        'CLOSE_WAIT': 0,
                        'LISTEN': 0
                    },
                    'interfaces': {}
                },
                'processes': {
                    'total': 0,
                    'top': []
                }
            }
        finally:
            self.stats_updating = False

    def get_process_info(self, pid):
        """获取指定进程的详细信息"""
        try:
            proc = psutil.Process(pid)
            info = {
                'pid': proc.pid,
                'name': proc.name(),
                'username': proc.username(),
                'status': proc.status(),
                'cpu_percent': proc.cpu_percent(interval=0.1),
                'memory_percent': proc.memory_percent(),
                'create_time': datetime.fromtimestamp(proc.create_time()).strftime('%Y-%m-%d %H:%M:%S'),
                'cmdline': ' '.join(proc.cmdline()),
                'num_threads': proc.num_threads(),
                'connections': len(proc.connections()),
                'open_files': len(proc.open_files()),
                'memory_info': {
                    'rss': proc.memory_info().rss,
                    'vms': proc.memory_info().vms
                },
                'io_counters': proc.io_counters()._asdict() if hasattr(proc, 'io_counters') else None,
                'nice': proc.nice(),
                'children': [child.pid for child in proc.children()]
            }
            return info
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess) as e:
            self.logger.error(f"获取进程信息失败 (PID: {pid}): {str(e)}")
            return None

    def kill_process(self, pid):
        """终止指定进程"""
        try:
            proc = psutil.Process(pid)
            proc.terminate()  # 先尝试正常终止
            try:
                proc.wait(timeout=3)  # 等待进程终止
            except psutil.TimeoutExpired:
                proc.kill()  # 如果超时则强制终止
            self.logger.info(f"已终止进程 (PID: {pid})")
            return True
        except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
            self.logger.error(f"终止进程失败 (PID: {pid}): {str(e)}")
            return False

    def get_process_tree(self):
        """获取进程树结构"""
        def get_children(parent_pid):
            children = []
            try:
                parent = psutil.Process(parent_pid)
                for child in parent.children(recursive=False):
                    try:
                        child_info = {
                            'pid': child.pid,
                            'name': child.name(),
                            'cpu_percent': child.cpu_percent(interval=0.1),
                            'memory_percent': child.memory_percent(),
                            'status': child.status(),
                            'children': get_children(child.pid)
                        }
                        children.append(child_info)
                    except (psutil.NoSuchProcess, psutil.AccessDenied):
                        continue
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass
            return children

        processes = []
        for proc in psutil.process_iter(['pid', 'ppid', 'name', 'username', 'cpu_percent', 'memory_percent']):
            try:
                pinfo = proc.info
                if pinfo['ppid'] == 1:  # 只获取顶级进程
                    pinfo['children'] = get_children(pinfo['pid'])
                    processes.append(pinfo)
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        return processes

    def get_history_data(self):
        """获取系统历史数据"""
        return self.history

    def update_history(self):
        """更新历史数据"""
        try:
            # 获取当前时间
            current_time = datetime.now().strftime('%H:%M:%S')
            
            # 获取当前系统状态
            stats = self.get_system_stats()
            if not stats:
                return
            
            # 计算网络速率
            if self.last_network_bytes:
                time_diff = time.time() - self.last_network_bytes['time']
                bytes_in_rate = (stats['network']['bytes_recv'] - self.last_network_bytes['recv']) / time_diff / 1024  # KB/s
                bytes_out_rate = (stats['network']['bytes_sent'] - self.last_network_bytes['sent']) / time_diff / 1024  # KB/s
            else:
                bytes_in_rate = 0
                bytes_out_rate = 0
            
            # 更新网络字节数记录
            self.last_network_bytes = {
                'time': time.time(),
                'recv': stats['network']['bytes_recv'],
                'sent': stats['network']['bytes_sent']
            }
            
            # 添加新的数据点
            self.history['timestamps'].append(current_time)
            self.history['cpu'].append(stats['cpu']['percent'])
            self.history['memory'].append(stats['memory']['percent'])
            self.history['network_in'].append(bytes_in_rate)
            self.history['network_out'].append(bytes_out_rate)
            self.history['connections'].append(len([c for c in psutil.net_connections() if c.status == 'ESTABLISHED']))
            self.history['disk_read'].append(stats['disk']['io']['read_bytes'])
            self.history['disk_write'].append(stats['disk']['io']['write_bytes'])
            
            # 限制历史数据点数量
            if len(self.history['timestamps']) > self.max_history_points:
                for key in self.history:
                    self.history[key] = self.history[key][-self.max_history_points:]
                
        except Exception as e:
            self.logger.error(f"更新历史数据失败: {str(e)}")

    def check_security(self):
        """执行安全检查"""
        try:
            self._check_suspicious_processes()
            self._check_network_attacks()
            self._check_file_integrity()
            self._check_system_anomalies()
        except Exception as e:
            self.logger.error(f"安全检查失败: {str(e)}")

    def _check_suspicious_processes(self):
        """检查可疑进程"""
        for proc in psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']):
            try:
                # 检查CPU和内存使用
                if proc.info['cpu_percent'] > self.security_settings['process_defense']['max_cpu_percent']:
                    self._handle_high_resource_usage(proc, 'CPU')
                
                if proc.info['memory_percent'] > self.security_settings['process_defense']['max_memory_percent']:
                    self._handle_high_resource_usage(proc, 'Memory')
                
                # 检查可疑进程名
                for suspicious in self.security_settings['process_defense']['suspicious_names']:
                    if suspicious in proc.info['name'].lower():
                        self._handle_suspicious_process(proc)
                        
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue

    def _check_network_attacks(self):
        """检查网络攻击"""
        connections = psutil.net_connections()
        ip_connections = {}
        
        for conn in connections:
            if conn.raddr and conn.raddr.ip:
                ip = conn.raddr.ip
                ip_connections[ip] = ip_connections.get(ip, 0) + 1
                
                # 检查连接数
                if ip_connections[ip] > self.security_settings['network_defense']['max_connections']:
                    self._handle_excessive_connections(ip)
                
                # 检查敏感端口
                if conn.laddr.port in self.security_settings['ip_defense']['suspicious_ports']:
                    self._monitor_sensitive_port(ip, conn.laddr.port)

    def _check_file_integrity(self):
        """检查文件完整性"""
        import hashlib
        
        for path in self.security_settings['file_defense']['protected_paths']:
            if os.path.exists(path):
                with open(path, 'rb') as f:
                    current_hash = hashlib.sha256(f.read()).hexdigest()
                    
                if path in self.file_checksums:
                    if current_hash != self.file_checksums[path]:
                        self._handle_file_modification(path)
                else:
                    self.file_checksums[path] = current_hash

    def _check_system_anomalies(self):
        """检查系统异常"""
        # 检查系统负载
        load1, load5, load15 = os.getloadavg()
        cpu_count = psutil.cpu_count()
        
        if load1 > cpu_count * 2:  # 负载超过CPU核心数的2倍
            self._handle_high_load(load1)
        
        # 检查磁盘使用
        disk = psutil.disk_usage('/')
        if disk.percent > 90:
            self._handle_disk_full(disk.percent)
        
        # 检查内存使用
        memory = psutil.virtual_memory()
        if memory.percent > 90:
            self._handle_memory_full(memory.percent)

    def _handle_high_resource_usage(self, proc, resource_type):
        """处理高资源使用"""
        message = f"进程资源使用过高 - {resource_type}: {proc.info['name']} (PID: {proc.info['pid']})"
        self.logger.warning(message)
        
        # 如果不是受保护进程，可以考虑终止
        if proc.info['name'] not in self.security_settings['process_defense']['protected_processes']:
            self.kill_process(proc.info['pid'])
            self.send_alert_email("资源使用告警", message)

    def _handle_suspicious_process(self, proc):
        """处理可疑进程"""
        message = f"发现可疑进程: {proc.info['name']} (PID: {proc.info['pid']})"
        self.logger.warning(message)
        self.kill_process(proc.info['pid'])
        self.send_alert_email("安全告警", message)

    def _handle_excessive_connections(self, ip):
        """处理过多连接"""
        message = f"IP连接数过多: {ip}"
        self.logger.warning(message)
        self.ban_ip(ip)
        self.send_alert_email("网络告警", message)

    def _handle_file_modification(self, path):
        """处理文件修改"""
        message = f"受保护文件被修改: {path}"
        self.logger.warning(message)
        self.send_alert_email("文件系统告警", message)

    def start(self):
        """启动防护服务"""
        self.logger.info("Linux防护服务已启动")
        
        # 启动Web界面
        from web_interface import WebInterface
        web = WebInterface(self)
        import threading
        web_thread = threading.Thread(target=web.start)
        web_thread.daemon = True
        web_thread.start()
        
        try:
            while True:
                self.check_connections()
                self.check_and_alert()
                self.update_history()
                self.check_security()  # 添加安全检查
                time.sleep(5)
        except KeyboardInterrupt:
            self.logger.info("服务已停止")

if __name__ == "__main__":
    defender = LinuxDefender()
    defender.start() 